package com.universest.swordholder.aitool

import com.alibaba.fastjson.JSONObject
import com.universest.swordholder.task.AbstractTask
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.combine
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.takeWhile
import kotlinx.coroutines.launch
import kotlinx.coroutines.suspendCancellableCoroutine
import okhttp3.sse.EventSource
import kotlin.coroutines.resume
import kotlin.coroutines.resumeWithException

class PromptTask(name: String = "",private val api: AbstractLlmApi) : AbstractTask<List<Message>,Message>(name) {

    private val _deepThought = MutableStateFlow("")

    val deepThought: Flow<String> = flow{
        combine(end,_deepThought) { end,deepThought ->
            Pair(end,deepThought)
        }.takeWhile { !it.first }
            .collect { emit(it.second) }
        emit(_deepThought.value)
    }

    private val _content = MutableStateFlow("")

    val content: Flow<String> = flow{
        combine(end,_content) { end,content ->
            Pair(end,content)
        }.takeWhile { !it.first }
            .collect { emit(it.second) }
        emit(_content.value)
    }

    @Volatile
    private var eventSource: EventSource? = null

    override suspend fun call(param: List<Message>): Message {
        _deepThought.emit("提示词长度：${param.getContextLength()}")
        val deepThoughtBuilder = StringBuilder()
        val contentBuilder = StringBuilder()
        return suspendCancellableCoroutine { continuation->
            eventSource = api.chatStream(param, object : AbstractLlmApi.StreamCallback {
                override fun onChunk(chunk: String) {
                    contentBuilder.append(chunk)
                    CoroutineScope(coroutineDispatcher).launch {
                        _content.emit(contentBuilder.toString())
                    }
                }

                override fun onDeepThought(reasoning: String) {
                    deepThoughtBuilder.append(reasoning)
                    CoroutineScope(coroutineDispatcher).launch {
                        _deepThought.emit(deepThoughtBuilder.toString())
                    }
                }

                override fun onComplete() {
                    continuation.resume(Message("assistant",contentBuilder.toString(),deepThoughtBuilder.toString()))
                }

                override fun onError(e: Exception) {
                    continuation.resumeWithException(e)
                }
            })
        }
    }

    override val progress: Flow<Pair<Int, Int>> = flow {
        var fakeProgress = 1.0
        combine(flow<Int> {
            while (!isEnd()) {
                fakeProgress = fakeProgress * (1 - Math.random() / 50)
                emit(((1 - fakeProgress) * 99).toInt())
                delay((Math.random()*1500).toLong())
            }
        },end) { countDown,end->
            Pair(countDown,end)
        }.takeWhile { !it.second }.collect { emit(it.first to 100) }
        emit(100 to 100)
    }

    override val message: Flow<String> = combine(deepThought,content) {
        deepThought,content ->
        val jsonObject = JSONObject()
        jsonObject.put("deepThought",deepThought)
        jsonObject.put("content",content)
        return@combine "aiProgress:${jsonObject.toJSONString()}"
    }

    override suspend fun cancel() {
        super.cancel()
        eventSource?.cancel()
    }
}